import os
import cv2
import utils
import numpy as np
from tqdm import tqdm
from random import randint


def rotate(ps, m):
    pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
    pts = np.hstack([pts, np.ones([len(pts), 1])]).T
    target_point = np.dot(m, pts)
    target_point = [[int(target_point[0][x]), int(target_point[1][x])] for x in range(len(target_point[0]))]
    return target_point


def rotate_img_and_point(img, points, angle, center_x, center_y, resize_rate=1.0):
    h, w, c = img.shape
    M = cv2.getRotationMatrix2D((center_x, center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    out_points = rotate(points, M)
    return res_img, out_points

def rotate_img(img, RotateRange, color=(0, 0, 0), resize_rate=1.0):
    h, w, c = img.shape
    angle = randint(RotateRange[0], RotateRange[1])
    M = cv2.getRotationMatrix2D((int(w / 2), int(h / 2)), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=color)
    return res_img

def CircleCoordinateRevise(OldCoordinate, NewCoordinate):
    SmallFlag = False
    X1old, Y1old, X2old, Y2old = OldCoordinate
    X1new, Y1new, X2new, Y2new = NewCoordinate

    Xmin = X1new + int((X2new - X1new - X2old + X1old) / 2)
    Ymin = Y1new + int((Y2new - Y1new - Y2old + Y1old) / 2)
    Xmax = X2new - int((X2new - X1new - X2old + X1old) / 2)
    Ymax = Y2new - int((Y2new - Y1new - Y2old + Y1old) / 2)

    # Xmin = max(Xmin, 0)
    # Ymin = max(Ymin, 0)
    if Xmax - Xmin < (X2old - X1old) * 0.8 or Ymax - Ymin < (Y2old - Y1old) * 0.8:
        SmallFlag = True
    return Xmin, Ymin, Xmax, Ymax, SmallFlag


def RotateImgXml(img, Infos, RotateRange):
    points = []
    for x1, y1, x2, y2, _ in Infos:
        points.append([x1, y1])
        points.append([x1, y2])
        points.append([x2, y1])
        points.append([x2, y2])
    NewImg, NewPoints = rotate_img_and_point(img, points, randint(RotateRange[0], RotateRange[1]),
                                             int(img.shape[1] / 2), int(img.shape[0] / 2))
    NewInfos = []
    for i in range(int(len(points) / 4)):
        ContinueFlag = False
        X1old, Y1old, X2old, Y2old = Infos[i][:-1]
        X1new, Y1new, X2new, Y2new = NewPoints[i * 4][0], NewPoints[i * 4 + 2][1], NewPoints[i * 4 + 3][0], \
                                     NewPoints[i * 4 + 1][1]
        # 圆形修正
        X1new, Y1new, X2new, Y2new, ContinueFlag = CircleCoordinateRevise([X1old, Y1old, X2old, Y2old],
                                                                          [X1new, Y1new, X2new, Y2new])
        if ContinueFlag:
            continue

        # 越界修正
        if X1new < 1:
            X1new = 1
        if Y1new < 1:
            Y1new = 1
        if X2new > img.shape[1]:
            X2new = img.shape[1]
        if Y2new > img.shape[0]:
            Y2new = img.shape[0]

        NewInfo = [X1new, Y1new, X2new, Y2new] + [Infos[i][-1]]
        NewInfos.append(NewInfo)
    return NewImg, NewInfos


if __name__ == '__main__':
    SrcPath = r"G:\Temp\Circle\Small"
    DstPath = r"G:\Temp\Circle\T\Temp"
    RotateRange = [-5, 5]
    Epochs = 1

    NameNum = 0

    for root, dirs, files in os.walk(SrcPath):
        for file in tqdm(files):
            if file[-1] == "g":
                for Epoch in range(Epochs):
                    img = cv2.imread(root + "//" + file)
                    Infos = utils.ReadXml(root + "//" + file[:-3] + "xml")
                    NewImg, NewInfos = RotateImgXml(img, Infos, RotateRange)
                    NewInfos.insert(0, DstPath + "\\R%05d.xml" % NameNum)
                    if len(NewInfos) > 1:
                        cv2.imwrite(NewInfos[0][:-3] + "jpg", NewImg)
                        utils.WriteXml(NewInfos, NewImg.shape[1], NewImg.shape[0])
                    NameNum += 1
